import torch
import torch.nn as nn
import torch.nn.functional as F

class HSwish(nn.Module):
    def forward(self, x):
        return x * F.relu6(x + 3) / 6

class SqueezeExcitation(nn.Module):
    def __init__(self, in_channels, reduction_ratio=4):
        super(SqueezeExcitation, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc1 = nn.Conv2d(in_channels, in_channels // reduction_ratio, kernel_size=1, bias=True)
        self.fc2 = nn.Conv2d(in_channels // reduction_ratio, in_channels, kernel_size=1, bias=True)
        
    def forward(self, x):
        out = self.avg_pool(x)
        out = self.fc1(out)
        out = F.relu(out)
        out = self.fc2(out)
        out = torch.sigmoid(out)
        out = out * x
        return out

class MobileNetV3Block(nn.Module):
    def __init__(self, in_channels, out_channels, expansion_factor, stride=1):
        super(MobileNetV3Block, self).__init__()
        self.expansion_factor = expansion_factor
        hidden_dim = int(in_channels * expansion_factor)
        
        self.conv1 = nn.Conv2d(in_channels, hidden_dim, kernel_size=1, stride=1, padding=0, bias=False)
        self.bn1 = nn.BatchNorm2d(hidden_dim)
        self.act1 = HSwish()
        
        self.conv2 = nn.Conv2d(hidden_dim, hidden_dim, kernel_size=3, stride=stride, padding=1, groups=hidden_dim, bias=False)
        self.bn2 = nn.BatchNorm2d(hidden_dim)
        self.act2 = HSwish()
        
        self.se = SqueezeExcitation(hidden_dim)
        
        self.conv3 = nn.Conv2d(hidden_dim, out_channels, kernel_size=1, stride=1, padding=0, bias=False)
        self.bn3 = nn.BatchNorm2d(out_channels)
        
        self.has_residual = stride == 1 and in_channels == out_channels
        
    def forward(self, x):
        identity = x
        
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.act1(out)
        
        out = self.conv2(out)
        out = self.bn2(out)
        out = self.act2(out)
        
        out = self.se(out)
        
        out = self.conv3(out)
        out = self.bn3(out)
        
        if self.has_residual:
            out += identity
        
        return out

class MobileNetV3Large(nn.Module):
    def __init__(self, num_classes=200):
        super(MobileNetV3Large, self).__init__()
        self.initial_conv = nn.Conv2d(3, 16, kernel_size=3, stride=2, padding=1, bias=False)
        self.bn = nn.BatchNorm2d(16)
        self.act = HSwish()
        
        self.blocks = nn.Sequential(
            MobileNetV3Block(16, 16, expansion_factor=1, stride=1),
            MobileNetV3Block(16, 24, expansion_factor=4, stride=2),
            MobileNetV3Block(24, 24, expansion_factor=3, stride=1),
            MobileNetV3Block(24, 40, expansion_factor=3, stride=2),
            MobileNetV3Block(40, 40, expansion_factor=3, stride=1),
            MobileNetV3Block(40, 40, expansion_factor=3, stride=1),
            MobileNetV3Block(40, 80, expansion_factor=6, stride=2),
            MobileNetV3Block(80, 80, expansion_factor=2.5, stride=1),
            MobileNetV3Block(80, 80, expansion_factor=2.3, stride=1),
            MobileNetV3Block(80, 80, expansion_factor=2.3, stride=1),
            MobileNetV3Block(80, 112, expansion_factor=6, stride=1),
            MobileNetV3Block(112, 112, expansion_factor=6, stride=1),
            MobileNetV3Block(112, 160, expansion_factor=6, stride=2),
            MobileNetV3Block(160, 160, expansion_factor=6, stride=1),
            MobileNetV3Block(160, 160, expansion_factor=6, stride=1),
            MobileNetV3Block(160, 320, expansion_factor=6, stride=1)
        )
        
        self.last_conv = nn.Conv2d(320, 1280, kernel_size=1, stride=1, padding=0, bias=False)
        self.bn_last = nn.BatchNorm2d(1280)
        self.act_last = HSwish()
        
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.dropout = nn.Dropout(p=0.2)
        self.fc = nn.Linear(1280, num_classes)
        
    def forward(self, x):
        out = self.initial_conv(x)
        out = self.bn(out)
        out = self.act(out)
        
        out = self.blocks(out)
        
        out = self.last_conv(out)
        out = self.bn_last(out)
        out = self.act_last(out)
        
        out = self.avgpool(out)
        out = self.dropout(out)
        out = out.view(out.size(0), -1)
        out = self.fc(out)
        
        return out
